"""
Token计数组件
负责计算输入和输出的token数量
"""
from typing import Dict, Any, Tuple, Optional
import logging
import re

# 设置日志
logger = logging.getLogger("connllm.tokens")

class TokenCounter:
	"""Token计数器，负责计算token数量和从响应中提取token信息"""
	
	def __init__(self, provider_type: str, model: str):
		"""
		初始化Token计数器
		
		Args:
			provider_type: 提供商类型
			model: 模型名称
		"""
		self.provider_type = provider_type
		self.model = model
		
	def count_tokens(self, text: str) -> int:
		"""
		计算文本的token数量
		
		Args:
			text: 要计算的文本
			
		Returns:
			token数量
		"""
		# 首先尝试使用提供商特定的tokenizer
		try:
			# 使用可选的tokenizers模块
			from ..optional.tokenizers import count_tokens
			return count_tokens(text, self.provider_type, self.model)
		except ImportError:
			pass
			
		# 其次尝试使用tiktoken
		try:
			import tiktoken
			
			# 选择编码器
			if 'gpt-4' in self.model or 'gpt-3.5' in self.model:
				encoding = tiktoken.encoding_for_model(self.model)
			else:
				encoding = tiktoken.get_encoding('cl100k_base')  # 默认编码器
				
			tokens = encoding.encode(text)
			return len(tokens)
		except (ImportError, Exception) as e:
			logger.debug(f"使用tiktoken计算token失败: {str(e)}")
			
		# 最后使用粗略估算
		return self._estimate_tokens(text)
		
	def _estimate_tokens(self, text: str) -> int:
		"""
		粗略估算token数量（仅作为后备方法）
		
		Args:
			text: 要估算的文本
			
		Returns:
			估算的token数量
		"""
		if not text:
			return 0
			
		# 英文文本估算（约0.25 tokens/char）
		if self._is_mainly_english(text):
			return max(1, len(text) // 4)
			
		# 中文文本估算（约1.0 tokens/char）
		if self._is_mainly_chinese(text):
			return len(text)
			
		# 其他文本估算（约0.33 tokens/char）
		return max(1, len(text) // 3)
		
	def _is_mainly_english(self, text: str) -> bool:
		"""判断文本是否主要为英文"""
		if not text:
			return True
			
		# 计算英文字符比例
		english_chars = sum(1 for c in text if c.isascii() and (c.isalnum() or c.isspace()))
		return english_chars / len(text) > 0.7
		
	def _is_mainly_chinese(self, text: str) -> bool:
		"""判断文本是否主要为中文"""
		if not text:
			return False
			
		# 计算中文字符比例（简繁中文）
		chinese_pattern = re.compile(r'[\u4e00-\u9fff\u3400-\u4dbf\U00020000-\U0002a6df\U0002a700-\U0002ebef]')
		chinese_chars = len(chinese_pattern.findall(text))
		return chinese_chars / len(text) > 0.5
		
	def extract_from_response(self, response: Any) -> Tuple[int, int]:
		"""
		从响应中提取token数量
		
		Args:
			response: API响应数据
			
		Returns:
			(输入token数, 输出token数)元组
		"""
		# 1. 尝试从可选模块导入函数
		try:
			from ..optional.tokenizers import extract_tokens_from_response
			
			# 提取token信息，包括缓存token
			input_tokens, output_tokens, _, _ = extract_tokens_from_response(
				response, self.provider_type, self.model
			)
			
			if input_tokens > 0 or output_tokens > 0:
				return input_tokens, output_tokens
		except ImportError:
			pass
			
		# 2. 使用特定提供商的响应格式提取token
		input_tokens, output_tokens = self._extract_from_provider_response(response)
		if input_tokens > 0 or output_tokens > 0:
			return input_tokens, output_tokens
			
		# 3. 如果无法从响应中直接获取，提取文本并估算
		input_text, output_text = self._extract_text_from_response(response)
		input_tokens = self.count_tokens(input_text) if input_text else 0
		output_tokens = self.count_tokens(output_text) if output_text else 0
		
		return input_tokens, output_tokens
		
	def _extract_from_provider_response(self, response: Any) -> Tuple[int, int]:
		"""
		从特定提供商的响应格式中提取token数量
		
		Args:
			response: API响应数据
			
		Returns:
			(输入token数, 输出token数)元组
		"""
		# 默认值
		input_tokens = 0
		output_tokens = 0
		
		# 只处理字典类型的响应
		if not isinstance(response, dict):
			return 0, 0
			
		# OpenAI风格
		if "usage" in response:
			usage = response["usage"]
			input_tokens = usage.get("prompt_tokens", 0)
			output_tokens = usage.get("completion_tokens", 0)
			if input_tokens > 0 or output_tokens > 0:
				return input_tokens, output_tokens
				
		# Anthropic风格
		if ("content" in response and isinstance(response.get("content"), list) and 
			response.get("usage") is not None):
			usage = response.get("usage", {})
			input_tokens = usage.get("input_tokens", 0)
			output_tokens = usage.get("output_tokens", 0)
			if input_tokens > 0 or output_tokens > 0:
				return input_tokens, output_tokens
		
		# 其他风格...
		
		return input_tokens, output_tokens
		
	def _extract_text_from_response(self, response: Any) -> Tuple[str, str]:
		"""
		从响应中提取输入和输出文本
		
		Args:
			response: API响应数据
			
		Returns:
			(输入文本, 输出文本)元组
		"""
		input_text = ""
		output_text = ""
		
		# 只处理字典类型的响应
		if not isinstance(response, dict):
			return input_text, output_text
			
		# 提取输入文本（如果有）
		if "prompt" in response:
			input_text = response["prompt"]
		elif "messages" in response:
			messages = response["messages"]
			if isinstance(messages, list):
				input_text = " ".join(msg.get("content", "") for msg in messages)
		
		# 提取输出文本（不同格式）
		# OpenAI风格
		if "choices" in response and len(response["choices"]) > 0:
			choice = response["choices"][0]
			if "message" in choice and "content" in choice["message"]:
				output_text = choice["message"]["content"]
			elif "text" in choice:
				output_text = choice["text"]
				
		# Anthropic风格
		elif "content" in response:
			if isinstance(response["content"], str):
				output_text = response["content"]
			elif isinstance(response["content"], list):
				output_text = " ".join(
					item.get("text", "") 
					for item in response["content"] 
					if item.get("type") == "text"
				)
				
		# Moonshot风格
		elif "response" in response:
			output_text = response["response"]
			
		return input_text, output_text


class CostCalculator:
	"""成本计算器，负责计算API调用成本"""
	
	def __init__(self, provider_type: str, model: str):
		"""
		初始化成本计算器
		
		Args:
			provider_type: 提供商类型
			model: 模型名称
		"""
		self.provider_type = provider_type
		self.model = model
		self.pricing_data = self._load_pricing_data()
		
	def _load_pricing_data(self) -> Dict[str, Any]:
		"""
		加载定价数据
		
		Returns:
			定价数据字典
		"""
		try:
			from ..core.pricing import get_pricing_data
			return get_pricing_data()
		except ImportError:
			logger.debug("无法导入定价数据模块，使用空定价数据")
			return {}
			
	def calculate(self, input_tokens: int, output_tokens: int) -> float:
		"""
		计算API调用成本
		
		Args:
			input_tokens: 输入token数量
			output_tokens: 输出token数量
			
		Returns:
			成本(美元)
		"""
		# 尝试使用标准定价函数
		try:
			from ..core.pricing import calculate_cost
			return calculate_cost(
				self.provider_type,
				self.model,
				input_tokens,
				output_tokens
			)
		except ImportError:
			# 如果标准函数不可用，使用内部实现
			return self._calculate_internal(input_tokens, output_tokens)
			
	def _calculate_internal(self, input_tokens: int, output_tokens: int) -> float:
		"""
		内部成本计算实现
		
		Args:
			input_tokens: 输入token数量
			output_tokens: 输出token数量
			
		Returns:
			成本(美元)
		"""
		# 本地提供商通常免费
		if self.provider_type in ['ollama', 'local']:
			return 0.0
			
		# 查找定价数据
		provider_prices = self.pricing_data.get(self.provider_type, {})
		model_prices = provider_prices.get(self.model, {})
		
		# 如果没有模型特定定价，尝试使用默认定价
		if not model_prices and self.model:
			# 尝试查找模型家族定价
			for model_pattern, prices in provider_prices.items():
				if model_pattern in self.model:
					model_prices = prices
					break
					
			# 尝试使用默认定价
			if not model_prices:
				model_prices = provider_prices.get("default", {})
		
		# 计算成本
		input_price = model_prices.get("input", 0.0) / 1000000  # 每百万token价格
		output_price = model_prices.get("output", 0.0) / 1000000
		
		cost = (input_tokens * input_price) + (output_tokens * output_price)
		return cost
